from langchain.prompts import PromptTemplate, ChatPromptTemplate
from langchain_community.llms import Tongyi
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableConfig
from typing import Literal, TypedDict
from langgraph.graph import StateGraph, START, END
from dotenv import load_dotenv, find_dotenv

load_dotenv(find_dotenv())

# Define the prompts we will route to
prompt_1 = PromptTemplate.from_template(
    '你是个优秀的物理学家，你很擅长用简明易懂的方式回答有关物理学的问题。'
    '请使用中文帮我解答下列问题：\n{input}'
)
prompt_2 = PromptTemplate.from_template(
    '你是个很好的数学家。你很擅长回答数学问题。'
    '请使用中文帮我解答下列问题：\n{input}'
)

llm = Tongyi()

# Construct the chains we will route to. These format the input query
# into the respective prompt, run it through a chat model, and cast
# the result to a string.
chain_1 = prompt_1 | llm | StrOutputParser()
chain_2 = prompt_2 | llm | StrOutputParser()

# Next: define the chain that selects which branch to route to.
# Here we will take advantage of tool-calling features to force
# the output to select one of two desired branches.
route_system = "Route the user's query to either the math or physics expert."
route_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", route_system),
        ("human", "{input}"),
    ]
)

# Define schema for output:
class RouteQuery(TypedDict):
    """Route query to destination expert."""
    destination: Literal["physics", "math"]
route_chain = route_prompt | llm.with_structured_output(RouteQuery)

# For LangGraph, we will define the state of the graph to hold the query,
# destination, and final answer.
class State(TypedDict):
    query: str
    destination: RouteQuery
    answer: str

# We define functions for each node, including routing the query:
async def route_query(state: State, config: RunnableConfig):
    destination = await route_chain.ainvoke(state["query"], config)
    return {"destination": destination}

# And one node for each prompt
async def prompt_1(state: State, config: RunnableConfig):
    return {"answer": await chain_1.ainvoke(state["query"], config)}


async def prompt_2(state: State, config: RunnableConfig):
    return {"answer": await chain_2.ainvoke(state["query"], config)}

# We then define logic that selects the prompt based on the classification
def select_node(state: State) -> Literal["prompt_1", "prompt_2"]:
    if state["destination"] == "physics":
        return "prompt_1"
    else:
        return "prompt_2"

# Finally, assemble the multi-prompt chain. This is a sequence of two steps:
# 1) Select "animal" or "vegetable" via the route_chain, and collect the answer
# alongside the input query.
# 2) Route the input query to chain_1 or chain_2, based on the
# selection.
graph = StateGraph(State)
graph.add_node("route_query", route_query)
graph.add_node("prompt_1", prompt_1)
graph.add_node("prompt_2", prompt_2)

graph.add_edge(START, "route_query")
graph.add_conditional_edges("route_query", select_node)
graph.add_edge("prompt_1", END)
graph.add_edge("prompt_2", END)
app = graph.compile()